package javaclass;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.NoSuchAlgorithmException;
import javax.crypto.*;
import javax.crypto.spec.*;
import sun.misc.BASE64Decoder;
import sun.misc.BASE64Encoder;

public class AES {

    public static SecretKeySpec getKeySpec() throws IOException, NoSuchAlgorithmException
    {
        // Get the KeyGenerator
        byte[] raw = new byte[16];
        File f = new File("aes_key");

        SecretKey key = null;
        SecretKeySpec skeySpec = null;

        if (f.exists()) {
            new FileInputStream(f).read(raw);
        } 
        else {
            KeyGenerator kgen = KeyGenerator.getInstance("AES");
            kgen.init(256);
            key = kgen.generateKey();
            raw = key.getEncoded();
            new FileOutputStream(f).write(raw);
        }

        skeySpec = new SecretKeySpec(raw,"AES");
        return skeySpec;
    }

    public String encrypt(String message) throws Exception{
        SecretKeySpec skeySpec = getKeySpec();
	Cipher cipher = Cipher.getInstance("AES");
        cipher.init(Cipher.ENCRYPT_MODE, skeySpec);
        BASE64Encoder enc = new BASE64Encoder();        
        String encryptedMessage = enc.encode(cipher.doFinal(message.getBytes()));
        return encryptedMessage;
    }

    public String decrypt(String message) throws Exception{
        SecretKeySpec skeySpec = getKeySpec();
	Cipher cipher = Cipher.getInstance("AES");
        cipher.init(Cipher.DECRYPT_MODE, skeySpec);
        BASE64Decoder dec = new BASE64Decoder();
        String originalString = new String(cipher.doFinal(dec.decodeBuffer(message)));
        return originalString;
    }
}
